import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
import torch

import matplotlib.pyplot as plt
import copy
from deepDeformOrt import *
import sys
# Neural Network
def transStr2Float(strIn, isGrad=0):
    floatList = []
    content = strIn.split(" ")
    for i in range(len(content)-1):
        tmp = float(content[i])
        floatList.append(tmp)
    floatList = torch.Tensor(floatList)
    if 1 == isGrad:
        floatList.requires_grad_(True)
    return floatList

def PlotDirCircle(pX, pY, pOrtX, pOrtY):
    r1 = 10
    for i in range(len(pX)):
        plt.scatter(pX[i], pY[i], c='r', marker='o', s = 8)
        X=[pX[i], pX[i] + r1 * pOrtX[i]]
        Y=[pY[i], (pY[i] + r1 * pOrtY[i])]
        plt.plot(X,Y, c='r', linestyle='-', linewidth = 1.0)

def Plot2DGrid(pX, pY, row, col):

    for i in range(row):
        pRowX = pX[i * col:(i + 1) * col]
        pRowY = pY[i * col:(i + 1) * col]
        for j in range(col-1):
            plt.plot([pRowX[j], pRowX[j + 1]], [pRowY[j], pRowY[j + 1]], c='b', linestyle='-', linewidth = 0.5)

    for i in range(col):
        for j in range(row-1):
            plt.plot([pX[j*col + i], pX[(j+1)*col + i]],[pY[j*col + i], pY[(j+1)*col + i]], c='b', linewidth = 0.5)

def Plot2DLine(pX1, pY1, pX2, pY2):
    n = pX1.size(0)
    for i in range(n):
        plt.plot([pX1[i], pX2[i]],[pY1[i], pY2[i]], c='g', linestyle='-', linewidth = 0.5)

def PlotDir2D(pointList1, pointList2, pGridList, strName):
    pX1 = pointList1[0]
    pY1 = pointList1[1]
    pOrtX1 = pointList1[2]
    pOrtY1 = pointList1[3]

    pX2 = pointList2[0]
    pY2 = pointList2[1]
    pOrtX2 = pointList2[2]
    pOrtY2 = pointList2[3]

    pGridx = pGridList[0]
    pGridy = pGridList[1]
    gridRow = pGridList[2]
    gridCol = pGridList[3]

    plt.figure(figsize=(6, 6), dpi=150)
    plt.gca().set_aspect(1)
    ax = plt.gca()
    ax.set_xlim(0, 600)  # 设置 x 轴最大最小刻度
    ax.set_ylim(0, 500)
    ax.invert_yaxis()
    PlotDirCircle(pX1, pY1, pOrtX1, pOrtY1)
    transX = 300
    PlotDirCircle(pX2.detach().numpy()+transX, pY2.detach().numpy(), pOrtX2.detach().numpy(), pOrtY2.detach().numpy())
    Plot2DGrid(pGridx.detach().numpy()+transX, pGridy.detach().numpy(), gridRow, gridCol)
    Plot2DLine(pX1, pY1, pX2.detach().numpy()+transX, pY2.detach().numpy())
    plt.savefig(strName)
    plt.show()

def TransPointCenter(pX1, pY1, pX2, pY2):
    with torch.no_grad():
        num = pX1.size(0)
        x1 = sum(pX1.squeeze(1))/num
        y1 = sum(pY1.squeeze(1))/num
        x2 = sum(pX2.squeeze(1))/num
        y2 = sum(pY2.squeeze(1))/num
        tx = x1 - x2
        ty = y1 - y2
    pX2 = pX2 + tx
    pY2 = pY2 + ty
    return tx, ty, pX2, pY2
def LoadData(dataName):
    dataList = []
    with open(dataName, 'r') as f:
        for line in f.readlines():
            print(line)
            dataList.append(line)

    content = dataList[0].split(" ")
    pointNum = int(content[0])

    pX1 = transStr2Float(dataList[1])
    pY1 = transStr2Float(dataList[2])
    pX1 = pX1.unsqueeze(1)
    pY1 = pY1.unsqueeze(1)
    pOrtX1 = transStr2Float(dataList[3])
    pOrtY1 = transStr2Float(dataList[4])

    pX2 = transStr2Float(dataList[5], isGrad=1)
    pY2 = transStr2Float(dataList[6], isGrad=1)
    pX2 = pX2.unsqueeze(1)
    pY2 = pY2.unsqueeze(1)
    pOrtX2 = transStr2Float(dataList[7])
    pOrtY2 = transStr2Float(dataList[8])

    content = dataList[13].split(" ")
    gridRow = int(content[0])
    gridCol = int(content[1])
    gridNum = gridRow * gridCol
    pGridx = transStr2Float(dataList[14])
    pGridy = transStr2Float(dataList[15])
    pGridx = pGridx.unsqueeze(1)
    pGridy = pGridy.unsqueeze(1)

    return pX1, pY1, pOrtX1, pOrtY1, pX2, pY2, pOrtX2, pOrtY2, pGridx, pGridy, gridRow, gridCol
def testArtificiallyPoints(lambda3):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dataName = 'Dir2DResult.txt'
    pX1, pY1, pOrtX1, pOrtY1, pX2, pY2, pOrtX2, pOrtY2, pGridx, pGridy, gridRow, gridCol = LoadData(dataName)
    #tx, ty, pX2, pY2 = TransPointCenter(pX1, pY1, pX2, pY2)
    PlotDir2D([pX1.squeeze(1), pY1.squeeze(1), pOrtX1, pOrtY1], [pX2.squeeze(1), pY2.squeeze(1), pOrtX2, pOrtY2], [pGridx.squeeze(1), pGridy.squeeze(1), gridRow, gridCol], "org.png")
    u,v = GetTransFunUV(pX1, pY1, pOrtX1, pOrtY1, pX2, pY2, pOrtX2, pOrtY2, lambda3, device)
    u = u.to('cpu')
    v = v.to('cpu')
    pTPointList2 = GetTransPoint(u, v, pX2, pY2, pOrtX2.unsqueeze(1), pOrtY2.unsqueeze(1))
    pTGridList = GetTransPoint(u, v, pGridx, pGridy, None, None)
    pTGridList.append(gridRow)
    pTGridList.append(gridCol)
    #画结果图
    name = "trans" + str(int(lambda3)) + ".png"
    PlotDir2D([pX1.squeeze(1), pY1.squeeze(1), pOrtX1, pOrtY1], pTPointList2, pTGridList, name) #"trans.png"

if __name__ == '__main__':
    print("传入参数的总长度为：", len(sys.argv))
    print("type:", type(sys.argv))
    print("function name:", sys.argv[0])
    try:
        print("第一个传入的参数为 lambda3:", sys.argv[1])
    except Exception as e:
        print("Input Error:", e)
    lambda3 = float(sys.argv[1])

    testArtificiallyPoints(lambda3)
    print("testArtificiallyPoints finish!")

